import pandas as pd
import re
import os

def CalMass(formula):
    mass_table = {
        'C': 12.0000,
        'H': 1.007825,
        'O': 15.994914,
        'N': 14.003074,
        'S': 31.9721,
        'P': 30.9738,
        'Cl': 34.9689,
        'Na': 22.9898,
        'K' : 38.9637,
        'Br': 78.9183,
        'F' : 18.9984,
        'I' : 126.9045,
        'D' : 2.0141
    }
    dic = {}
    for key in mass_table.keys():
        all_item_of_key = re.findall(key, formula)
        if all_item_of_key != []:
            parttern = re.compile(f'{key}[\d]+|{key}[a-z]+|{key}')
            num_str = re.findall(parttern, formula)
            num = 0
            for s in num_str:
                len_dic = s.split(key)
                if len_dic[1].isdigit():
                    num += int(len_dic[1])
                elif len_dic[1] == '':
                    num += 1
            dic[key] = num
        else:
            dic[key] = 0
    mass = 0
    for key in mass_table.keys():
        mass += (mass_table[key]*dic[key])
    return mass

def calculate_adduct_masses(excel_file, sheet_name=0, formula_col='分子式', output_file=None):

    df = pd.read_excel(excel_file, sheet_name=sheet_name)
    
    adduct_increments = {
        'M+H': 1.00783,
        'M+Na': 22.98977,
        'M+K': 38.9637,
        'M+NH4': 18.03382,
        'M-H': -1.00783,
        'M-H2O+H': -18.010565 + 1.00783,
        'M+FA-H': 45.01755
    }
    
    df['分子量'] = df[formula_col].apply(CalMass)
    
    for adduct, increment in adduct_increments.items():
        df[adduct] = increment
    
    if output_file is None:
        base, ext = os.path.splitext(excel_file)
        output_file = f"{base}_adductsdiff{ext}"
    
    df.to_excel(output_file, index=False)
    print(f"结果已保存到 {output_file}")
    return df

if __name__ == "__main__":
    input_file = "./Data/杏贝混标.xlsx"  
    calculate_adduct_masses(input_file)